package com.feidee.transformer;

import com.feidee.fdspark.transformer.ColAssembler;
import org.dmg.pmml.*;
import org.jpmml.converter.ContinuousFeature;
import org.jpmml.converter.Feature;
import org.jpmml.converter.PMMLUtil;
import org.jpmml.sparkml.FeatureConverter;
import org.jpmml.sparkml.SparkMLEncoder;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

/**
 * @Author songhaicheng
 * @Date 2018/12/19 15:23
 * @Description
 * @Reviewer
 */
public class ColAssemblerConverter extends FeatureConverter<ColAssembler> {

    public static final int MIN_ARITY = 2;

    public ColAssemblerConverter(ColAssembler transformer) {
        super(transformer);
    }

    @Override
    public List<Feature> encodeFeatures(SparkMLEncoder encoder) {
        List<Expression> expressions = new ArrayList<>();
        Constant delimiter = PMMLUtil.createConstant(getTransformer().getDelimiter(), DataType.STRING);
        // 先保留原来的特征列
        String[] cols = getTransformer().getCols();
        // 当拼接列参数只有一个时，拼接上一个空字符串，防止报 concat function 最少实参长度（2）错误，同时又保证值正确
        if (cols.length < MIN_ARITY) {
            Feature colFeature = encoder.getOnlyFeature(cols[0]);
            expressions.add(colFeature.ref());
            expressions.add(PMMLUtil.createConstant("", DataType.STRING));
        } else {
            for (int i = 0; i < cols.length; i++) {
                Feature colFeature = encoder.getOnlyFeature(cols[i]);
                // 用分隔符拼接特征列
                expressions.add(colFeature.ref());
                if (i != cols.length - 1) {
                    expressions.add(delimiter);
                }
            }
        }
        // 再保存拼接后的结果
        Apply concat = PMMLUtil.createApply("concat", expressions.toArray(new Expression[0]));
        DerivedField output = encoder.createDerivedField(FieldName.create(getTransformer().getOutputCol()),
                OpType.CATEGORICAL,
                DataType.STRING,
                concat);
        return Collections.singletonList(new ContinuousFeature(encoder, output));
    }
}
